#include "swtiger.h"
#include "utils.h"

using namespace std;
using namespace UTILS;

SWTIGER::SWTIGER()
{
    NumActions = 3;
    NumObservations = 2;
    RewardRange = 100;
    Discount = 0.95;
    InitGeneral();
}

void SWTIGER::InitGeneral()
{
    RandomSeed(0);
}

STATE* SWTIGER::Copy(const STATE& state) const
{
    const SWTIGER_STATE& rockstate = safe_cast<const SWTIGER_STATE&>(state);
    SWTIGER_STATE* newstate = MemoryPool.Allocate();
    *newstate = rockstate;
    return newstate;
}

void SWTIGER::Validate(const STATE& state) const
{
    const SWTIGER_STATE& rockstate = safe_cast<const SWTIGER_STATE&>(state);
}

STATE* SWTIGER::CreateStartState() const
{
    SWTIGER_STATE* rockstate = MemoryPool.Allocate();
    //get a possible start state
    if (Bernoulli(0.5))
    	rockstate->state = 0;
    else
    	rockstate->state = 1;
    return rockstate;
}

void SWTIGER::FreeState(STATE* state) const
{
    SWTIGER_STATE* rockstate = safe_cast<SWTIGER_STATE*>(state);
    MemoryPool.Free(rockstate);
}

bool SWTIGER::Step(STATE& state, int action,
    int& observation, double& reward) const
{
	// if terminal, return true

    SWTIGER_STATE& rockstate = safe_cast<SWTIGER_STATE&>(state);
    // listen
    if (action == 0) {
    	reward = -1;

        if (Bernoulli(0.85))
        	observation = rockstate.state;
        else
        	observation = 1-rockstate.state;

    }
    // open-left
    if (action == 1) {
    	// state is left
    	if (rockstate.state == 0) reward = 10;
    	else reward = -100;

        if (Bernoulli(0.5))
        	rockstate.state = 0;
        else
        	rockstate.state = 1;

        if (Bernoulli(0.5))
        	observation = 0;
        else
        	observation = 1;
       // return true;

    }
    // open-right
    if (action == 2) {
    	// state is left
    	if (rockstate.state == 0) reward = -100;
    	else reward = 10;

        if (Bernoulli(0.5))
        	rockstate.state = 0;
        else
        	rockstate.state = 1;

        if (Bernoulli(0.5))
        	observation = 0;
        else
        	observation = 1;
        //return true;
    }

    return false;
}


void SWTIGER::DisplayBeliefs(const BELIEF_STATE& beliefState,
    std::ostream& ostr) const
{
	//ostr << "Belief" << endl;
	//ostr << "Number of Samples" << beliefState.GetNumSamples() << endl;
	int numZeroes = 0;
	int numOnes = 0;
	for (int i=0; i < beliefState.GetNumSamples(); i++) {
		const SWTIGER_STATE* tigerstate = safe_cast<const SWTIGER_STATE*>(beliefState.GetSample(i));
		//ostr << tigerstate->state << endl;
		if (tigerstate->state == 0) numZeroes++;
		if (tigerstate->state == 1) numOnes++;
	}
	ostr << "Belief" << beliefState.GetNumSamples() << ":" << (double) numZeroes/beliefState.GetNumSamples() << "(" << numZeroes << ")" \
			<< "," << (double) numOnes/beliefState.GetNumSamples() << "(" << numOnes << ")" << endl;
}

void SWTIGER::DisplayState(const STATE& state, std::ostream& ostr) const
{
    const SWTIGER_STATE& rockstate = safe_cast<const SWTIGER_STATE&>(state);
    ostr << "State:" << rockstate.state;
    ostr << endl;
}

void SWTIGER::DisplayObservation(const STATE& state, int observation, std::ostream& ostr) const
{
    ostr << "Observation:" << observation << endl;
}

void SWTIGER::DisplayAction(int action, std::ostream& ostr) const
{
	ostr << "Action:" << action << endl;
}
